#!/usr/bin/env python3
import os
import sys
import grpc
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))

import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections
from p4runtime_lib.error_utils import printGrpcError

def install_ipv4_rule(p4info_helper, switch, dst_ip, dst_mac, out_port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip, 32)},
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": out_port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"[{switch.name}] IPv4 Rule: {dst_ip} -> {dst_mac} via port {out_port}")

def install_ipv6_rule(p4info_helper, switch, dst_ip, dst_mac, out_port, prefix_len=64):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv6_lpm",
        match_fields={"hdr.ipv6.dstAddr": (dst_ip, prefix_len)},
        action_name="MyIngress.ipv6_forward",
        action_params={
            "dstAddr": dst_mac,
            "port": out_port
        }
    )
    switch.WriteTableEntry(table_entry)
    print(f"[{switch.name}] IPv6 Rule: {dst_ip}/{prefix_len} -> {dst_mac} via port {out_port}")

def main():
    # 更新为level3的编译输出文件
    p4info_path = "build/level3.p4.p4info.txt"
    bmv2_json_path = "build/level3.json"
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_path)

    try:
        # 初始化交换机连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-level3-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-level3-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-level3-requests.txt')
        s4 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s4',
            address='127.0.0.1:50054',
            device_id=3,
            proto_dump_file='logs/s4-level3-requests.txt')

        # 配置流水线
        for sw in [s1, s2, s3, s4]:
            sw.MasterArbitrationUpdate()
            sw.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_json_path)
            print(f"[INFO] Pipeline configured on {sw.name}")

        # ========== IPv4核心路径规则 h1→S1→S2→S4→h2 ==========
        # S1规则 (h1→s2)
        install_ipv4_rule(p4info_helper, s1, "10.0.4.4", "08:00:00:00:02:00", 2)  # h2流量强制走s2
        install_ipv4_rule(p4info_helper, s1, "10.0.1.1", "08:00:00:00:01:11", 1)  # 返回流量
        install_ipv4_rule(p4info_helper, s1, "10.0.1.10", "08:00:00:00:01:00", 1)

        # S2规则 (s1→s4)
        install_ipv4_rule(p4info_helper, s2, "10.0.4.4", "08:00:00:00:04:00", 2)  # 转发到s4
        install_ipv4_rule(p4info_helper, s2, "10.0.1.1", "08:00:00:00:01:00", 1)  # 返回流量到s1

        # S4规则 (s2→h2)
        install_ipv4_rule(p4info_helper, s4, "10.0.4.4", "08:00:00:00:04:44", 3)  # 转发到h2
        install_ipv4_rule(p4info_helper, s4, "10.0.1.1", "08:00:00:00:02:00", 1)  # 返回流量到s2

        # ========== IPv4备用路径规则 ==========
        # S1的备用路径（到s3）
        install_ipv4_rule(p4info_helper, s1, "10.0.3.30", "08:00:00:00:03:00", 3)  # 其他流量

        # S3规则（虽然实验不要求，但保证互通）
        install_ipv4_rule(p4info_helper, s3, "10.0.4.4", "08:00:00:00:04:00", 2)  # 到h2
        install_ipv4_rule(p4info_helper, s3, "10.0.1.1", "08:00:00:00:01:00", 1)  # 返回h1

        # S4的备用路径（来自s3）
        install_ipv4_rule(p4info_helper, s4, "10.0.1.1", "08:00:00:00:01:00", 1)  # 来自s3的返回流量

        # ========== IPv6核心路径规则 h1→S1→S3→S4→h2 ==========
        # 定义IPv6地址常量
        H1_IPV6 = "2001:db8::1"
        H2_IPV6 = "2001:db8::4"
        S1_GW_IPV6 = "2001:db8::a"  # s1作为h1的IPv6网关
        S4_GW_IPV6 = "2001:db8::b"  # s4作为h2的IPv6网关

        # S1规则 (h1→s3)
        install_ipv6_rule(p4info_helper, s1, H2_IPV6, "08:00:00:00:03:00", 3)  # h2流量强制走s3
        install_ipv6_rule(p4info_helper, s1, H1_IPV6, "08:00:00:00:01:11", 1)  # 返回流量

        # S3规则 (s1→s4) - 只支持IPv6
        install_ipv6_rule(p4info_helper, s3, H2_IPV6, "08:00:00:00:04:00", 2)  # 转发到s4
        install_ipv6_rule(p4info_helper, s3, H1_IPV6, "08:00:00:00:01:00", 1)  # 返回流量到s1

        # S4规则 (s3→h2)
        install_ipv6_rule(p4info_helper, s4, H2_IPV6, "08:00:00:00:04:44", 3)  # 转发到h2
        install_ipv6_rule(p4info_helper, s4, H1_IPV6, "08:00:00:00:03:00", 2)  # 返回流量到s3

        # 添加IPv6网关路由规则
        install_ipv6_rule(p4info_helper, s1, S1_GW_IPV6, "08:00:00:00:01:00", 1)  # s1的IPv6网关
        install_ipv6_rule(p4info_helper, s4, S4_GW_IPV6, "08:00:00:00:04:00", 3)  # s4的IPv6网关

        # s1上，加允许ff02::1:ff00:444多播转发到s4方向（h2在s4）
        install_ipv6_rule(p4info_helper, s1, "ff02::1:ff00:444", "08:00:00:00:03:00", 3)
        # s4上，加允许ff02::1:ff00:444多播转发到h2端口
        install_ipv6_rule(p4info_helper, s4, "ff02::1:ff00:444", "08:00:00:00:04:44", 3)
        # h1链路本地相关（在s1，转发到h1）
        install_ipv6_rule(p4info_helper, s1, "fe80::a00:ff:fe00:111", "08:00:00:00:01:11", 1)
        # h2链路本地相关（在s4，转发到h2）
        install_ipv6_rule(p4info_helper, s4, "fe80::a00:ff:fe00:444", "08:00:00:00:04:44", 3)

        # NDP多播（h1和h2各自的ff02::1:ff00:xxxx）在所有hop都加
        install_ipv6_rule(p4info_helper, s1, "ff02::1:ff00:444", "08:00:00:00:03:00", 3)
        install_ipv6_rule(p4info_helper, s3, "ff02::1:ff00:444", "08:00:00:00:04:00", 2)
        install_ipv6_rule(p4info_helper, s4, "ff02::1:ff00:444", "08:00:00:00:04:44", 3)
        # 反方向同理
        install_ipv6_rule(p4info_helper, s4, "ff02::1:ff00:111", "08:00:00:00:03:00", 2)
        install_ipv6_rule(p4info_helper, s3, "ff02::1:ff00:111", "08:00:00:00:01:00", 1)
        install_ipv6_rule(p4info_helper, s1, "ff02::1:ff00:111", "08:00:00:00:01:11", 1)

        print("\n[SUCCESS] All rules installed!")
        print("IPv4 Path: h1→s1→s2→s4→h2")
        print("IPv6 Path: h1→s1→s3→s4→h2")

    except grpc.RpcError as e:
        printGrpcError(e)
    finally:
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    main()
